import numpy as np
import time
from scipy.spatial.distance import cdist
from ldp_searching import LDP_Searching
from lmst_clu_opt import LMSTCLU_OPT

def LDPMST_OPT(A, clu_num):
    """
    基于最小生成树的局部密度峰聚类 (LDP-MST)
    输入:
        A: ndarray, shape (N, dim)
        clu_num: int — 期望簇数
    输出:
        cl: ndarray, shape (N,) — 最终聚类标签
        elapsed: float — 运行时长（秒）
    """
    A = np.asarray(A)
    N, _ = A.shape
    print('LDP-Searching...')
    start = time.time()

    index, supk, nb, rho, local_core, cores2, cl_init, cluster_number = LDP_Searching(A)
    print(f'初始子簇个数为：{cluster_number}')

    # 核点距离矩阵
    core_pts = A[cores2]
    core_dist = cdist(core_pts, core_pts)
    maxd = core_dist.max()

    # 构建每个核点的扩展集合
    cdataexp = []
    nc = []
    for core_idx in cores2:
        members = np.where(local_core == core_idx)[0]
        nc.append(len(members))
        neigh = index[members, 1:supk+1].flatten()
        exp_set = np.union1d(members, neigh)
        cdataexp.append(exp_set)

    # 调整核点间距离
    for i in range(cluster_number):
        for j in range(i+1, cluster_number):
            inset1 = np.intersect1d(cdataexp[i], cdataexp[j])
            if inset1.size == 0:
                core_dist[i, j] = maxd * (core_dist[i, j] + 1)
            else:
                averho = rho[inset1].sum()
                core_dist[i, j] = core_dist[i, j] / (averho * inset1.size)
            core_dist[j, i] = core_dist[i, j]

    minsize = int(0.018 * N)
    print('LDP-MST构建...')
    cores_cl = LMSTCLU_OPT(core_pts, core_dist, nc, minsize, clu_num) # 在局部密度峰值上构建 MST 并做子簇分裂 (LMSTCLU_OPT)

    # TODO 生成最终标签
    cl = np.zeros(N, dtype=int)
    for idx, core_idx in enumerate(cores2): # idx：核心点的序号，core_idx：核心点在数据中的索引
        cl[core_idx] = cores_cl[idx]
    for i in range(N):
        if local_core[i] >= 0:
            # 如果点 i 是某个核心点的邻居，则将其标签设置为该核心点的标签
            cl[i] = cl[local_core[i]]
        else:
            # 否则，将其标签设置为 0
            cl[i] = 0
    elapsed = time.time() - start
    print(f'运行时长: {elapsed:.3f} s')
    print('---LDP-MST Complete---')

    ncluster = cl.max() # 簇的数量
    return A, cl, ncluster